package com.dc.ai.model;

import com.dc.ai.security.utils.SecurityUtils;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.data.redis.core.RedisTemplate;

import java.time.Duration;
import java.util.List;

public class RedisChatMemory implements ChatMemory {

    private final RedisTemplate<String, Message> redisTemplate;
    private static final String KEY_PREFIX = "chat:history:";
    private static final String MEMORY_KEY_PREFIX = "chat:memory:";

    public RedisChatMemory(RedisTemplate<String, Message> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        String key = getKey(conversationId);
        // 批量追加消息到列表尾部
        redisTemplate.opsForList().rightPushAll(key, messages);
        // 设置过期时间（示例：7天）
        redisTemplate.expire(key, Duration.ofDays(7));
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        String key = getKey(conversationId);
        // 获取列表总长度
        Long size = redisTemplate.opsForList().size(key);
        if (size == null || size == 0) {
            return List.of();
        }
        // 计算起始索引（支持lastN大于总长度的情况）
        long start = Math.max(0, size - lastN);
        return redisTemplate.opsForList().range(key, start, -1);
    }

    @Override
    public void clear(String conversationId) {
        redisTemplate.delete(getKey(conversationId));
    }

    @Override
    public void add(String id, Message message) {
        // 获取当前用户ID并添加到消息元数据中
        String userId = SecurityUtils.getUserId();
        if (userId != null && !userId.isEmpty()) {
            message.getMetadata().put("userId", userId);
        }
        
        // 将消息添加到Redis列表
        redisTemplate.opsForList().rightPush(MEMORY_KEY_PREFIX + id, message);
    }

    private String getKey(String conversationId) {
        return KEY_PREFIX + conversationId;
    }
}